import torch
import pandas as pd
from torch.utils.data import Dataset


class RadarDataset(Dataset):
    def __init__(self, df: pd.DataFrame, seq, device, target_cols):
        super(RadarDataset, self).__init__()
        self.seq = seq
        source_cols = [item for item in df.columns if item not in target_cols]
        self.x = df[source_cols].to_numpy()
        self.y = df[target_cols].to_numpy()
        self.x = torch.tensor(self.x, dtype=torch.float32).to(device)
        self.y = torch.tensor(self.y, dtype=torch.float32).to(device)
        
    def __getitem__(self, index):
        return self.x[index: index + self.seq], self.y[index + self.seq]

    def __len__(self):
        return len(self.x) - self.seq
        

if __name__ == '__main__':
    from load_data import loadradar
    df = loadradar()
    radardataset = RadarDataset(df, target_cols=['RadWdSpd1'])